package com.test.langchain4j.config;

import com.test.langchain4j.constant.QdrantConstant;
import com.test.langchain4j.service.ChatAssistant;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.store.embedding.qdrant.QdrantEmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

/**
 * Created with IntelliJ IDEA.
 *
 * @description:
 * @author: liuziyang
 * @since: 2025/8/4 11:15
 * @modifiedBy:
 * @version: 1.0
 */
@Configuration
public class LLMConfig {
  @Bean
  public ChatModel chatModel() {
    return OpenAiChatModel.builder()
        .apiKey(System.getenv("qwen-api-key"))
        .modelName("qwen-plus")
        .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1")
        .build();
  }

  /**
   * 需要预处理文档并将其存储在专门的嵌入存储（也称为矢量数据库）中。当用户提出问题时，这对于快速找到相关信息是必要的。 我们可以使用我们支持的 15
   * 多个嵌入存储中的任何一个，但为了简单起见，我们将使用内存中的嵌入存储：
   *
   * <p>https://docs.langchain4j.dev/integrations/embedding-stores/in-memory
   *
   * @return
   */
  @Primary
  @Bean
  public InMemoryEmbeddingStore<TextSegment> embeddingStore() {
    return new InMemoryEmbeddingStore<>();
  }

  @Bean
  public QdrantClient qdrantClient() {
    QdrantGrpcClient.Builder grpcClientBuilder =
        QdrantGrpcClient.newBuilder("10.2.54.170", 6334, false);
    return new QdrantClient(grpcClientBuilder.build());
  }

  @Bean
  public EmbeddingStore<TextSegment> qdrantEmbeddingStore() {
    return QdrantEmbeddingStore.builder()
        .host("10.2.54.170")
        .port(6334)
        .collectionName(QdrantConstant.COLLECTION_NAME)
        .build();
  }

  @Bean
  public ChatAssistant assistant(ChatModel chatModel, EmbeddingStore<TextSegment> embeddingStore) {
    return AiServices.builder(ChatAssistant.class)
        .chatModel(chatModel)
        .chatMemory(MessageWindowChatMemory.withMaxMessages(50))
        .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))
        .build();
  }
}
